/*
 * @Author: dadadaXU 1413107032@qq.com
 * @Date: 2025-01-01 14:57:29
 * @LastEditors: dadadaXU 1413107032@qq.com
 * @LastEditTime: 2025-01-05 14:56:22
 * @FilePath: /FileBackup/src/fileCompress.cpp
 * @Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
 */
#include "../include/fileCompress.hpp"

FileCompressor::FileCompressor(std::filesystem::path source_file_path, std::filesystem::path target_file_path)
    : sourceFilePath(std::filesystem::absolute(source_file_path).lexically_normal()),
      targetFilePath(std::filesystem::absolute(target_file_path).lexically_normal())
{
    if (!std::filesystem::exists(source_file_path))
    {
        std::cerr << "Invalid path: " << source_file_path << std::endl;
        return;
    }

    std::memset(tokenTable, 0, sizeof(tokenTable));
    codeTable.clear();
    huffmanTree = nullptr;
}

FileCompressor::~FileCompressor()
{
    DeleteHuffmanTree(this->huffmanTree);
}

void FileCompressor::CountToken(FileBaseUnit &source_file)
{
    unsigned char buffer[BUFFER_SIZE] = {0};
    size_t buffer_count = 0;

    source_file.seekg(0, std::ios::beg);
    do
    {
        source_file.read((char *)(buffer), BUFFER_SIZE);
        buffer_count = source_file.gcount();
        for (size_t i = 0; i < buffer_count; i++)
            this->tokenTable[buffer[i]]++;
    } while (buffer_count > 0 && source_file.peek() != EOF);

    if (source_file.fail() && !source_file.eof())
    {
        std::cerr << "File read error: " << source_file.filePath << std::endl;
        return;
    }
}

void FileCompressor::WriteToken(FileBaseUnit &target_file)
{
    /* 写指针跳过文件头 */
    target_file.seekp(std::streamoff(sizeof(CompressHeader)), std::ios::beg);

    target_file.write((const char *)tokenTable, sizeof(tokenTable));
    target_file.fileHeader.CRC ^=
        target_file.Calc_CRC32((uint8_t *)(tokenTable), sizeof(tokenTable));
}

void FileCompressor::ReadToken(FileBaseUnit &source_file)
{
    /* 读指针跳过文件头 */
    source_file.seekg(std::streamoff(sizeof(CompressHeader)), std::ios::beg);

    source_file.read((char *)tokenTable, sizeof(tokenTable));
    source_file.fileHeader.CRC ^=
        source_file.Calc_CRC32((uint8_t *)(tokenTable), sizeof(tokenTable));
}

void FileCompressor::CreateHuffmanTree()
{
    /* 构建词频小根堆 */
    for (size_t i = 0; i < TOKEN_NUM; i++)
    {
        if (tokenTable[i] > 0) // 只包含存在的词
        {
            HuffmanTNode *node = new HuffmanTNode;
            node->token = static_cast<unsigned char>(i);
            node->frequency = tokenTable[i];
            node->lchild = nullptr;
            node->rchild = nullptr;
            frequencyHeap.push(node);
        }
    }

    if (frequencyHeap.size() == 0)
    {
        std::cerr << "TokenTable is empty!" << std::endl;
        return;
    }

    /* 自底向上构建 Huffman 树 */
    while (frequencyHeap.size() > 1)
    {
        HuffmanTNode *left = frequencyHeap.top();
        frequencyHeap.pop();
        HuffmanTNode *right = frequencyHeap.top();
        frequencyHeap.pop();

        HuffmanTNode *root = new HuffmanTNode;
        root->lchild = left;
        root->rchild = right;
        root->frequency = left->frequency + right->frequency;
        frequencyHeap.push(root);
    }
    huffmanTree = frequencyHeap.top();
    frequencyHeap.pop();
}

void FileCompressor::DeleteHuffmanTree(HuffmanTNode *root)
{
    if (root == nullptr)
        return;

    if (root->lchild != nullptr)
        DeleteHuffmanTree(root->lchild);
    if (root->rchild != nullptr)
        DeleteHuffmanTree(root->rchild);

    delete root;
    root = nullptr;
}

void FileCompressor::HuffmanEncode(HuffmanTNode *root, std::string code)
{
    root->code = code;
    if (root->lchild == nullptr && root->rchild == nullptr)
    {
        codeTable[root->token] = code;
        return;
    }

    if (root->lchild != nullptr)
        HuffmanEncode(root->lchild, code + "0");
    if (root->rchild != nullptr)
        HuffmanEncode(root->rchild, code + "1");
}

bool FileCompressor::EncodeFileData(FileBaseUnit &source_file, FileBaseUnit &target_file)
{
    unsigned char buffer_in[BUFFER_SIZE] = {0};
    size_t buffer_in_count = 0;

    std::string code_str;
    unsigned char bit_code = 0; // 压缩后的比特码
    size_t bit_count = 0;

    /* 重置读指针 */
    source_file.clear(); // 清除结束标志
    source_file.seekg(0, std::ios::beg);
    do
    {
        source_file.read((char *)(buffer_in), BUFFER_SIZE);
        buffer_in_count = source_file.gcount();
        for (size_t i = 0; i < buffer_in_count; i++)
        {
            /* 数据编码 */
            code_str = codeTable[buffer_in[i]];
            /* 编码压缩 */
            for (auto c : code_str)
            {
                bit_code = (bit_code << 1) | (c - '0'); // 低位插入 c
                bit_count++;
                if (bit_count % 8 == 0) // 写入数据
                {
                    target_file.put(bit_code);
                    bit_count = 0;
                }
            }
        }
        /* 计算原数据的校验码 */
        target_file.fileHeader.CRC ^=
            target_file.Calc_CRC32((uint8_t *)(buffer_in), buffer_in_count);
    } while (buffer_in_count > 0 && source_file.peek() != EOF);

    if (bit_count > 0)
    {
        bit_code <<= (8 - (bit_count % 8));
        target_file.put(bit_code);
    }

    if (source_file.fail() && !source_file.eof())
    {
        std::cerr << "File read error: " << source_file.filePath << std::endl;
        return false;
    }

    return true;
}

bool FileCompressor::DecodeFileData(FileBaseUnit &source_file, FileBaseUnit &target_file, off_t file_size)
{
    unsigned char buffer_in[BUFFER_SIZE] = {0};
    unsigned char buffer_out[BUFFER_SIZE] = {0};
    size_t buffer_in_count = 0;
    size_t buffer_out_count = 0;

    unsigned char bit_code = 0; // 压缩后的比特码
    size_t bit_count = 0;

    HuffmanTNode *p = this->huffmanTree; // 访问 Huffman 树指针
    off_t decompress_size = 0;           // 当前解压程度
    // TODO 文件解压进度条 decompress_size / file_size

    do
    {
        source_file.read((char *)(buffer_in), BUFFER_SIZE);
        buffer_in_count = source_file.gcount();
        for (size_t i = 0; i < buffer_in_count; i++)
        {
            bit_code = buffer_in[i];
            bit_count = 0;
            /* 访问 Huffman 树 */
            while (bit_count < 8)
            {
                /* 从左往右解码 */
                if (bit_code & 0x80)
                    p = p->rchild; // 右1
                else
                    p = p->lchild; // 左0

                if (p->lchild == nullptr && p->rchild == nullptr)
                {
                    target_file.put(p->token);
                    /* 校验码计算（保持算法一致）*/
                    buffer_out[buffer_out_count++] = p->token;
                    if (buffer_out_count == BUFFER_SIZE)
                    {
                        source_file.fileHeader.CRC ^=
                            source_file.Calc_CRC32((uint8_t *)(buffer_out), BUFFER_SIZE);
                        buffer_out_count = 0;
                    }

                    decompress_size++;
                    if (decompress_size >= file_size) // 解压成功
                    {
                        if (buffer_out_count > 0)
                            source_file.fileHeader.CRC ^=
                                source_file.Calc_CRC32((uint8_t *)(buffer_out), buffer_out_count);
                        return true;
                    }
                    p = this->huffmanTree;
                }
                bit_count++;
                bit_code <<= 1;
            }
        }
    } while (buffer_in_count > 0 && source_file.peek() != EOF);

    if (source_file.fail() && !source_file.eof())
    {
        std::cerr << "File read error: " << source_file.filePath << std::endl;
        return false;
    }

    return false;
}

bool FileCompressor::Compress()
{
    /* 待压缩文件流 */
    FileBaseUnit compressing_file_in(sourceFilePath);
    if (!(compressing_file_in.GetFileType() & (FILE_TYPE_NORMAL | FILE_TYPE_HARD_LINK)))
    {
        std::cerr << "Unsupport file type - " << sourceFilePath << std::endl;
        return false;
    }

    /* 压缩文件路径 */
    if (!std::filesystem::exists(targetFilePath.parent_path()))
        std::filesystem::create_directories(targetFilePath.parent_path());
    if (std::filesystem::is_directory(targetFilePath)) // 路径为目录则默认文件名
        targetFilePath.append(DEFAULT_COMPRESS_FILENAME);
    if (targetFilePath.extension() != FILE_SUFFIX_COMPRESS)
        targetFilePath += FILE_SUFFIX_COMPRESS; // 追加后缀名

    /* 创建压缩文件 */
    CompressHeader compress_header = {0};
    FileBaseUnit compressFile_out(targetFilePath, compress_header);

    /* 构建 Huffman 树 */
    CountToken(compressing_file_in);
    CreateHuffmanTree();
    HuffmanEncode(this->huffmanTree, "");
    /* 词频（根结点 = 文件大小）*/
    if (this->huffmanTree->frequency != compressing_file_in.fileHeader.metadata.st_size)
    {
        std::cerr << "Create Huffman tree failed: " << this->huffmanTree->frequency << std::endl;
        return false;
    }

    WriteToken(compressFile_out);
    if (!EncodeFileData(compressing_file_in, compressFile_out))
    {
        std::cerr << "Encode file data failed: " << sourceFilePath << std::endl;
        return false;
    }

    // TODO 计算压缩比
    /* 写入文件头 */
    std::strncpy(compress_header.filePath, sourceFilePath.c_str(), TAR_FILE_PATH);
    compress_header.CRC = compressFile_out.fileHeader.CRC;
    compress_header.offset = sizeof(CompressHeader);
    compress_header.fileSize = compressing_file_in.fileHeader.metadata.st_size;
    compressFile_out.UpdateCompressHeader(compress_header);

    return true;
}

bool FileCompressor::Decompress(bool go_back)
{
    if (sourceFilePath.extension() != FILE_SUFFIX_COMPRESS)
    {
        std::cerr << "Invalid file: " << sourceFilePath << std::endl;
        return false;
    }

    /* 读取压缩文件头 */
    FileBaseUnit decompressFile_in(sourceFilePath);
    CompressHeader compress_header = {0};
    if (!decompressFile_in.ReadCompressHeader(compress_header))
    {
        std::cerr << "Read file header failed: " << sourceFilePath << std::endl;
        return false;
    }

    /* 切换工作路径 */
    if (go_back) // 文件原始路径
    {
        targetFilePath = std::filesystem::path(compress_header.filePath).parent_path();
        if (!std::filesystem::exists(targetFilePath))
        {
            std::cerr << "Reload file path failed: " << targetFilePath << std::endl;
            return false;
        }
    }
    if (!std::filesystem::exists(targetFilePath.parent_path()))
        std::filesystem::create_directories(targetFilePath.parent_path());

    /* 还原文件后缀名 */
    targetFilePath.replace_extension(
        std::filesystem::path(compress_header.filePath).extension());

    /* 构造解压文件 */
    FileHeader fileHeader = {0};
    fileHeader.metadata.st_mode |= S_IFREG; // 构造普通文件
    std::strncpy(fileHeader.filePath, targetFilePath.c_str(), COMPRESS_FILE_PATH);
    FileBaseUnit decompressing_file_out(fileHeader);

    /* 构建 Huffman 树 */
    ReadToken(decompressFile_in);
    CreateHuffmanTree();
    if (this->huffmanTree->frequency != compress_header.fileSize)
    {
        std::cerr << "Recreate Huffman tree failed: " << this->huffmanTree->frequency << std::endl;
        return false;
    }

    if (!DecodeFileData(decompressFile_in, decompressing_file_out, compress_header.fileSize))
    {
        std::cerr << "Decode file data failed: " << sourceFilePath << std::endl;
        return false;
    }

    /* 校验码检查 */
    if (compress_header.CRC != decompressFile_in.fileHeader.CRC)
    {
        std::cerr << "Decompress file data failed! " << std::endl
                  << "\tOriginal CRC:   " << compress_header.CRC << std::endl
                  << "\tDecompress CRC: " << decompressFile_in.fileHeader.CRC << std::endl;
        return false;
    }

    return true;
}